#include <string.h>
#include <stdlib.h>

const char* dgemm_desc = "Simple blocked dgemm.";

#if !defined(BIG_BLOCK_SIZE)
#define BIG_BLOCK_SIZE 440
#define SMALL_BLOCK_SIZE 88 
#define ALIGN_SIZE 64
#endif

#define min(a,b) (((a)<(b))?(a):(b))


/* This auxiliary subroutine performs a smaller dgemm operation
 *  C := C + A * B
 * where C is M-by-N, A is M-by-K, and B is K-by-N.
 * C and B are column major, A is row major
 */
static void inline do_block (int n, int M, int N, int K, 
        double* restrict A, double* restrict B, double* restrict C)
{
        
    /* For each column j of B */ 
    for (int j=0; j < N; j++) {

        
        /* For each row i of A */

        /* Loop unrolling */
        int i=0;

        for (i = 0; i < M-3; i+=4) {
            /* Compute C(i,j) */
            double cij = C[i+j*n];
            double ci1j = C[i+1+j*n];
            double ci2j = C[i+2+j*n];
            double ci3j = C[i+3+j*n];
            for (int k = 0; k < K; ++k) {
                cij += A[i*K + k] * B[j*K + k];
                ci1j += A[(i+1)*K + k] * B[j*K + k];
                ci2j += A[(i+2)*K + k] * B[j*K + k];
                ci3j += A[(i+3)*K + k] * B[j*K + k];
            }
            C[i+j*n] = cij;
            C[i+1+j*n] = ci1j;
            C[i+2+j*n] = ci2j;
            C[i+3+j*n] = ci3j;
        }
    
        for (; i < M; ++i) {
            /* Compute C(i,j) */
            double cij = C[i+j*n];
            for (int k = 0; k < K; ++k)
                cij += A[i*K + k] * B[j*K + k];
            C[i+j*n] = cij;
        }
    }
}


/* This auxiliary subroutine performs a smaller (yet big enough) dgemm operation
 *  C := C + A * B
 * where C is M-by-N, A is M-by-K, and B is K-by-N. 
 * C and B are column major, A is row major
 */
static void inline do_big_block (int n, int M_, int N_, int K_, 
        double* restrict A, double* restrict B, double* restrict C)
{
    int t=0;

    /* For each block-row of A */ 
    for (int i = 0; i < M_; i += SMALL_BLOCK_SIZE) {

        int M = min (SMALL_BLOCK_SIZE, M_-i);

        /* For each block-column of B */
        for (int  j = 0; j < N_; j += SMALL_BLOCK_SIZE) {

            int N = min (SMALL_BLOCK_SIZE, N_-j);

            //double* restrict Ctemp;
            //t = posix_memalign((void**)&Ctemp, ALIGN_SIZE, sizeof(double)*M*N);
            //if (t!=0)   exit(1);


            ///* Copy stuff from Original C to our little local buffer */
            //for (int x=0; x<N; x++) 
            //    memcpy(Ctemp+(x*M), C+((x+j)*M_)+i, M*sizeof(double));


            /* Accumulate block dgemms into block of C */
            for (int k = 0; k < K_; k += SMALL_BLOCK_SIZE) {
                /* Correct block dimensions if block "goes off edge of" the matrix */

                int K = min (SMALL_BLOCK_SIZE, K_-k);

                /* Perform individual block dgemm */

                double* restrict Atemp;
                double* restrict Btemp;
            
                t = posix_memalign((void**)&Atemp, ALIGN_SIZE, sizeof(double)*M*K);
                if (t!=0) exit(1);
                t = posix_memalign((void**)&Btemp, ALIGN_SIZE, sizeof(double)*K*N);
                if (t!=0) exit(1);


                /* Let's copy A to Atemp, make that row-major */
                for (int x=0; x<M; x++) {
                    memcpy (Atemp + (x*K), A+((x+i)*K_)+k, K*sizeof(double));
                }

                /* Let's copy B to Btemp, make that column-major */
                for (int x=0; x<N; x++) {
                    memcpy (Btemp +(x*K), B+((x+j)*K_)+k, K*sizeof(double));
                }


                // do_block(n, M, N, K, A + i + k*n, B + k + j*n, C + i + j*n);

                /* Naive matrix multiplication
                   C = C + AxB
                   Just that the matrices are no longer all nxn
                 */
                do_block(n, M, N, K, Atemp, Btemp, C + (i+ j*n));

                free(Atemp);
                free(Btemp);
            }

            ///* Copy back */
            //for (int x=0; x<N; x++) 
            //    memcpy(C+((x+j)*M_)+i, Ctemp+(x*M), M*sizeof(double));

            //free(Ctemp);
        }
    }
}

/* This routine performs a dgemm operation
 *  C := C + A * B
 * where A, B, and C are n-by-n matrices stored in column-major format. 
 * On exit, A and B maintain their input values. */  
void square_dgemm (int n, double* A, double* B, double* restrict C)
{


    int t=0;

    /* There shouldn't be any harm done in making A row major right in the
     * beginning. Let's do that to make our life simpler.
     */

    double* restrict A_;
    t = posix_memalign((void **)&A_, ALIGN_SIZE, sizeof(double)*n*n);
    if (t!=0)   exit(1);

    for (int i=0; i<n; i++) {
        for (int j=0; j<n; j++) {
            A_[i*n+j] = A[i+n*j];
        }
    }

    /* For each block-row of A_ */ 
    for (int i = 0; i < n; i += BIG_BLOCK_SIZE) {

        int M = min (BIG_BLOCK_SIZE, n-i);

        /* For each block-column of B */
        for (int  j = 0; j < n; j += BIG_BLOCK_SIZE) {

            int N = min (BIG_BLOCK_SIZE, n-j);

            //double* restrict Ctemp;
            //t = posix_memalign((void**)&Ctemp, ALIGN_SIZE, sizeof(double)*M*N);
            //if (t!=0)   exit(1);


            ///* Copy stuff from Original C to our little local buffer */
            //for (int x=0; x<N; x++) 
            //    memcpy(Ctemp+(x*M), C+((x+j)*n)+i, M*sizeof(double));


            /* Accumulate block dgemms into block of C */
            for (int k = 0; k < n; k += BIG_BLOCK_SIZE) {
                /* Correct block dimensions if block "goes off edge of" the matrix */

                int K = min (BIG_BLOCK_SIZE, n-k);

                /* Perform individual block dgemm */

                double* restrict Atemp;
                double* restrict Btemp;
            
                t = posix_memalign((void**)&Atemp, ALIGN_SIZE, sizeof(double)*M*K);
                if (t!=0) exit(1);
                t = posix_memalign((void**)&Btemp, ALIGN_SIZE, sizeof(double)*K*N);
                if (t!=0) exit(1);


                /* Let's copy A_ to Atemp, make that row-major */
                for (int x=0; x<M; x++) {
                    memcpy (Atemp + (x*K), A_+((x+i)*n)+k, K*sizeof(double));
                }

                /* Let's copy B to Btemp, make that column-major */
                for (int x=0; x<N; x++) {
                    memcpy (Btemp +(x*K), B+((x+j)*n)+k, K*sizeof(double));
                }


                /* Naive matrix multiplication
                   C = C + AxB
                   Just that the matrices are no longer all nxn
                 */
                do_big_block(n, M, N, K, Atemp, Btemp, C + (i + j*n));

                free(Atemp);
                free(Btemp);
            }

            ///* Copy back */
            //for (int x=0; x<N; x++) 
            //    memcpy(C+((x+j)*n)+i, Ctemp+(x*M), M*sizeof(double));

            //free(Ctemp);
        }
    }
}
